import requests
import json
import base64
import os
import time
from datetime import datetime

def encode_image_to_base64(image_path):
    """将图片文件编码为base64字符串"""
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')

def decode_base64_to_image(base64_string, output_path):
    """将base64字符串解码并保存为图片文件"""
    image_data = base64.b64decode(base64_string)
    with open(output_path, "wb") as image_file:
        image_file.write(image_data)
        
def get_mime_type(file_path):
    """根据文件扩展名获取MIME类型"""
    ext = os.path.splitext(file_path)[1].lower()
    mime_types = {
        '.png': 'image/png',
        '.jpg': 'image/jpeg',
        '.jpeg': 'image/jpeg',
        '.webp': 'image/webp',
        '.gif': 'image/gif'
    }
    return mime_types.get(ext, 'image/png')  # 默认为PNG

def generate_image_with_gemini(input_image_path, second_image_path, prompt="照片一中女生穿上牛仔裤和图二中靴子, 保持女生的姿势不变, 面部特征表情、背景和整体真实感不变。镜头拉远保证女生全身和靴子均在图内"):
    """使用Gemini API生成图片"""
    
    # API配置
    url = "https://yunwu.ai/v1beta/models/gemini-2.5-flash-image-preview:generateContent?key=sk-8Wa80xngQs8U1zfmEBHE0TQQIJ02cNEI76Yc4FkTUDvYrBes"
    
    # 读取并编码两张输入图片
    print(f"正在读取第一张图片: {input_image_path}")
    image1_base64 = encode_image_to_base64(input_image_path)
    
    print(f"正在读取第二张图片: {second_image_path}")
    image2_base64 = encode_image_to_base64(second_image_path)
    
    # 获取图片的MIME类型
    mime_type1 = get_mime_type(input_image_path)
    mime_type2 = get_mime_type(second_image_path)
    
    # 构建请求数据
    payload = json.dumps({
        "contents": [
            {
                "role": "user",
                "parts": [
                    {
                        "text": prompt
                    },
                    {
                        "inline_data": {
                            "mime_type": mime_type1,
                            "data": image1_base64
                        }
                    },
                    {
                        "inline_data": {
                            "mime_type": mime_type2,
                            "data": image2_base64
                        }
                    }
                ]
            }
        ],
        "generationConfig": {
            "responseModalities": [
                "IMAGE"
            ]
        }
    })
    
    headers = {
        'Content-Type': 'application/json'
    }
    
    # 发送请求（带重试机制）
    print("正在调用Gemini API...")
    max_retries = 3
    retry_delay = 5  # 重试间隔5秒
    
    for attempt in range(max_retries):
        try:
            print(f"第 {attempt + 1} 次尝试调用API...")
            response = requests.post(url, headers=headers, data=payload, timeout=60)
            
            if response.status_code == 200:
                result = response.json()
                print("API调用成功！")
                print("API响应结构:", list(result.keys()) if result else "空响应")
                
                # 提取生成的图片数据
                if 'candidates' in result and len(result['candidates']) > 0:
                    candidate = result['candidates'][0]
                    print("找到候选结果")
                    print("候选结果结构:", list(candidate.keys()) if candidate else "空候选")
                    
                    if 'content' in candidate and 'parts' in candidate['content']:
                        parts = candidate['content']['parts']
                        print(f"找到 {len(parts)} 个部分")
                        
                        for i, part in enumerate(parts):
                            print(f"部分 {i} 结构:", list(part.keys()) if part else "空部分")
                            
                            # 检查两种可能的数据结构
                            inline_data = None
                            if 'inline_data' in part:
                                inline_data = part['inline_data']
                            elif 'inlineData' in part:
                                inline_data = part['inlineData']
                            
                            if inline_data:
                                mime_type = inline_data.get('mime_type') or inline_data.get('mimeType', 'unknown')
                                print(f"部分 {i} 类型: {mime_type}")
                                
                                if mime_type == 'image/png' or mime_type.startswith('image/'):
                                    # 生成输出文件名 - 基于第一张图片的文件名
                                    base_name = os.path.splitext(os.path.basename(input_image_path))[0]
                                    output_filename = f"{base_name}_ai.png"
                                    output_path = os.path.join("pic", output_filename)
                                    
                                    # 解码并保存图片
                                    print(f"正在保存生成的图片: {output_path}")
                                    decode_base64_to_image(inline_data['data'], output_path)
                                    print(f"图片已保存到: {output_path}")
                                    return output_path
                                else:
                                    print(f"部分 {i} 不是PNG图片，跳过")
                            else:
                                print(f"部分 {i} 没有inline_data或inlineData")
                
                print("未找到生成的图片数据")
                print("完整API响应:", json.dumps(result, indent=2, ensure_ascii=False)[:1000] + "..." if len(str(result)) > 1000 else json.dumps(result, indent=2, ensure_ascii=False))
                return None
                
            elif response.status_code == 429:  # 请求过于频繁
                print(f"API请求过于频繁 (429)，等待 {retry_delay * 2} 秒后重试...")
                if attempt < max_retries - 1:
                    time.sleep(retry_delay * 2)
                    continue
            else:
                print(f"API调用失败，状态码: {response.status_code}")
                print("错误信息:", response.text)
                if attempt < max_retries - 1:
                    print(f"等待 {retry_delay} 秒后重试...")
                    time.sleep(retry_delay)
                    continue
                    
        except requests.exceptions.ConnectionError as e:
            print(f"连接错误 (第 {attempt + 1} 次尝试): {e}")
            if "ProxyError" in str(e) or "10054" in str(e):
                print("检测到代理连接问题或连接被重置")
            if attempt < max_retries - 1:
                print(f"等待 {retry_delay} 秒后重试...")
                time.sleep(retry_delay)
                continue
        except requests.exceptions.Timeout as e:
            print(f"请求超时 (第 {attempt + 1} 次尝试): {e}")
            if attempt < max_retries - 1:
                print(f"等待 {retry_delay} 秒后重试...")
                time.sleep(retry_delay)
                continue
        except requests.exceptions.RequestException as e:
            print(f"请求异常 (第 {attempt + 1} 次尝试): {e}")
            if attempt < max_retries - 1:
                print(f"等待 {retry_delay} 秒后重试...")
                time.sleep(retry_delay)
                continue
        except Exception as e:
            print(f"未知错误 (第 {attempt + 1} 次尝试): {e}")
            if attempt < max_retries - 1:
                print(f"等待 {retry_delay} 秒后重试...")
                time.sleep(retry_delay)
                continue
    
    print(f"API调用失败，已重试 {max_retries} 次")
    return None

def generate_image_with_gemini_three_images(input_image_path, second_image_path, third_image_path, prompt="请基于这三张图片生成一张新的图片"):
    """使用Gemini API生成图片（处理3张输入图片）"""
    
    # API配置
    url = "https://yunwu.ai/v1beta/models/gemini-2.5-flash-image-preview:generateContent?key=sk-8Wa80xngQs8U1zfmEBHE0TQQIJ02cNEI76Yc4FkTUDvYrBes"
    
    # 读取并编码三张输入图片
    print(f"正在读取第一张图片: {input_image_path}")
    image1_base64 = encode_image_to_base64(input_image_path)
    
    print(f"正在读取第二张图片: {second_image_path}")
    image2_base64 = encode_image_to_base64(second_image_path)
    
    print(f"正在读取第三张图片: {third_image_path}")
    image3_base64 = encode_image_to_base64(third_image_path)
    
    # 获取图片的MIME类型
    mime_type1 = get_mime_type(input_image_path)
    mime_type2 = get_mime_type(second_image_path)
    mime_type3 = get_mime_type(third_image_path)
    
    # 构建请求数据
    payload = json.dumps({
        "contents": [
            {
                "role": "user",
                "parts": [
                    {
                        "text": prompt
                    },
                    {
                        "inline_data": {
                            "mime_type": mime_type1,
                            "data": image1_base64
                        }
                    },
                    {
                        "inline_data": {
                            "mime_type": mime_type2,
                            "data": image2_base64
                        }
                    },
                    {
                        "inline_data": {
                            "mime_type": mime_type3,
                            "data": image3_base64
                        }
                    }
                ]
            }
        ],
        "generationConfig": {
            "responseModalities": [
                "IMAGE"
            ]
        }
    })
    
    headers = {
        'Content-Type': 'application/json'
    }
    
    # 发送请求
    print("正在调用Gemini API...")
    response = requests.post(url, headers=headers, data=payload)
    
    if response.status_code == 200:
        result = response.json()
        print("API调用成功！")
        print("API响应结构:", list(result.keys()) if result else "空响应")
        
        # 提取生成的图片数据
        if 'candidates' in result and len(result['candidates']) > 0:
            candidate = result['candidates'][0]
            print("找到候选结果")
            print("候选结果结构:", list(candidate.keys()) if candidate else "空候选")
            
            if 'content' in candidate and 'parts' in candidate['content']:
                parts = candidate['content']['parts']
                print(f"找到 {len(parts)} 个部分")
                
                for i, part in enumerate(parts):
                    print(f"部分 {i} 结构:", list(part.keys()) if part else "空部分")
                    
                    # 检查两种可能的数据结构
                    inline_data = None
                    if 'inline_data' in part:
                        inline_data = part['inline_data']
                    elif 'inlineData' in part:
                        inline_data = part['inlineData']
                    
                    if inline_data:
                        mime_type = inline_data.get('mime_type') or inline_data.get('mimeType', 'unknown')
                        print(f"部分 {i} 类型: {mime_type}")
                        
                        if mime_type == 'image/png' or mime_type.startswith('image/'):
                            # 生成输出文件名 - 基于第一张图片的文件名
                            base_name = os.path.splitext(os.path.basename(input_image_path))[0]
                            output_filename = f"{base_name}_ai_pose.png"
                            output_path = os.path.join("pic", output_filename)
                            
                            # 解码并保存图片
                            print(f"正在保存生成的图片: {output_path}")
                            decode_base64_to_image(inline_data['data'], output_path)
                            print(f"图片已保存到: {output_path}")
                            return output_path
                        else:
                            print(f"部分 {i} 不是PNG图片，跳过")
                    else:
                        print(f"部分 {i} 没有inline_data或inlineData")
        
        print("未找到生成的图片数据")
        print("完整API响应:", json.dumps(result, indent=2, ensure_ascii=False)[:1000] + "..." if len(str(result)) > 1000 else json.dumps(result, indent=2, ensure_ascii=False))
    else:
        print(f"API调用失败，状态码: {response.status_code}")
        print("错误信息:", response.text)
    
    return None

def process_images_with_boots(girl_image_path, boots_image_path, prompt="生成一张图片： 照片一中女生穿上牛仔裤和图二中靴子, 保持女生的姿势不变, 面部特征表情、背景和整体真实感不变。广角镜头保证女生全身和靴子均在图内。"):
    """
    处理女生图片和靴子图片，生成新的合成图片
    
    Args:
        girl_image_path (str): 女生图片路径
        boots_image_path (str): 靴子图片路径
        prompt (str): 提示词，默认为"图中女生穿上牛仔裤和图二中靴子, 保持脸部不变, 镜头拉远保证女生全身和靴子均在图内"
    
    Returns:
        str: 生成的图片路径，如果失败返回None
    """
    # 检查输入文件是否存在
    if not os.path.exists(girl_image_path):
        print(f"错误: 女生图片文件不存在: {girl_image_path}")
        return None
    
    if not os.path.exists(boots_image_path):
        print(f"错误: 靴子图片文件不存在: {boots_image_path}")
        return None
    
    # 生成图片
    generated_image = generate_image_with_gemini(
        girl_image_path, 
        boots_image_path,
        prompt
    )
    
    if generated_image:
        print(f"图片生成完成！输出文件: {generated_image}")
        return generated_image
    else:
        print("图片生成失败")
        return None

def process_images_with_boots_pose(girl_image_path, boots_image_path, pose_image_path, prompt="图中女生穿上牛仔裤和图二中靴子, 身体比例尽量如图三所示, 保持脸部不变, 镜头拉远保证女生全身和靴子均在图内"):
    """
    处理女生图片、靴子图片和姿势参考图片，生成新的合成图片
    
    Args:
        girl_image_path (str): 女生图片路径
        boots_image_path (str): 靴子图片路径
        pose_image_path (str): 姿势参考图片路径
        prompt (str): 提示词，默认为"图中女生穿上牛仔裤和图二中靴子, 身体比例尽量如图三所示, 保持脸部不变, 镜头拉远保证女生全身和靴子均在图内"
    
    Returns:
        str: 生成的图片路径，如果失败返回None
    """
    # 检查输入文件是否存在
    if not os.path.exists(girl_image_path):
        print(f"错误: 女生图片文件不存在: {girl_image_path}")
        return None
    
    if not os.path.exists(boots_image_path):
        print(f"错误: 靴子图片文件不存在: {boots_image_path}")
        return None
    
    if not os.path.exists(pose_image_path):
        print(f"错误: 姿势参考图片文件不存在: {pose_image_path}")
        return None
    
    # 生成图片
    generated_image = generate_image_with_gemini_three_images(
        girl_image_path, 
        boots_image_path,
        pose_image_path,
        prompt
    )
    
    if generated_image:
        print(f"图片生成完成！输出文件: {generated_image}")
        return generated_image
    else:
        print("图片生成失败")
        return None
if __name__ == "__main__":
    # 示例使用
    girl_image = "pic/2023-04-25_18.45.08_白飞飞_清华日常_春天的照片_2.png"
    boots_image = "src/boots.png"
    pose_image = "src/pose.png"
    
    # 使用默认提示词处理2张图片
    result = process_images_with_boots(girl_image, boots_image)
    
    # 使用默认提示词处理3张图片（包含姿势参考）
    # result_pose = process_images_with_boots_pose(girl_image, boots_image, pose_image)
    
    # 或者使用自定义提示词
    # custom_prompt = "让女生穿上图二的靴子，保持原图风格"
    # result = process_images_with_boots(girl_image, boots_image, custom_prompt)
    
    # custom_prompt_pose = "让女生穿上图二的靴子，参考图三的姿势，保持原图风格"
    # result_pose = process_images_with_boots_pose(girl_image, boots_image, pose_image, custom_prompt_pose)
